"""
将训练好的Keras模型转换为CoreML格式（用于iOS）
"""

import coremltools as ct
from tensorflow import keras
import numpy as np

def rebuild_model(input_shape: tuple = (1, 117)) -> keras.Model:
    """
    重建模型架构（兼容TensorFlow 2.12+）
    必须与train_pipeline.py中的build_model()完全一致
    
    根据config.yaml:
    - lstm_layers: 2
    - lstm_units: 64
    - dropout: 0.3
    - dense_units: 32
    
    Args:
        input_shape: (time_steps, features)
    
    Returns:
        重建的Keras模型
    """
    model = keras.Sequential([
        keras.layers.Input(shape=input_shape),
        
        # 第1层LSTM（return_sequences=True因为有第2层）
        keras.layers.LSTM(64, return_sequences=True, dropout=0.3),
        
        # 第2层LSTM
        keras.layers.LSTM(64, return_sequences=False, dropout=0.3),
        
        # 全连接层
        keras.layers.Dense(32, activation='relu'),
        keras.layers.Dropout(0.3),
        
        # 输出层
        keras.layers.Dense(1, activation='sigmoid')
    ])
    
    # 编译（必须编译才能加载权重）
    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=0.001),
        loss='binary_crossentropy',
        metrics=['accuracy', keras.metrics.Precision(), keras.metrics.Recall()]
    )
    
    return model

def convert_to_coreml(
    keras_model_path: str = "best_model.h5",  # 当前目录
    output_path: str = "ChewingDetector.mlmodel",  # 当前目录
    input_shape: tuple = (1, 117)  # (time_steps=1, features=117)
):
    """
    转换Keras模型到CoreML
    
    Args:
        keras_model_path: Keras模型路径
        output_path: 输出CoreML模型路径
        input_shape: 输入形状 (time_steps, feature_dim)
    """
    print(f"加载Keras模型: {keras_model_path}")
    
    # 方法1: 尝试直接加载
    try:
        model = keras.models.load_model(keras_model_path)
        print("✓ 直接加载成功")
    except (TypeError, ValueError) as e:
        print(f"直接加载失败: {e}")
        print("使用兼容方法: 重建模型架构并加载权重...")
        
        # 方法2: 重建模型并加载权重
        model = rebuild_model(input_shape=input_shape)
        model.load_weights(keras_model_path)
        print("✓ 权重加载成功")
    
    print(f"模型输入形状: {model.input_shape}")
    print(f"模型输出形状: {model.output_shape}")
    
    # 转换为CoreML（使用neuralnetwork格式，更稳定且易于修改）
    print("\n开始转换到CoreML...")
    
    coreml_model = ct.convert(
        model,
        convert_to="neuralnetwork"  # 使用经典格式（更稳定）
    )
    
    # 获取spec并修改
    spec = coreml_model.get_spec()
    
    # 获取原始输入输出名称
    old_input_name = spec.description.input[0].name
    old_output_name = spec.description.output[0].name
    print(f"原始输入名称: {old_input_name}")
    print(f"原始输出名称: {old_output_name}")
    
    # 重命名输入和输出
    ct.utils.rename_feature(spec, old_input_name, "input_features")
    ct.utils.rename_feature(spec, old_output_name, "chewing_probability")
    print(f"✓ 已重命名为: input_features -> chewing_probability")
    
    # 添加元数据
    spec.description.metadata.author = "Chewing Detection Training Pipeline"
    spec.description.metadata.shortDescription = "实时咀嚼检测模型"
    spec.description.metadata.versionString = "1.0"
    
    # 添加输入输出描述
    spec.description.input[0].shortDescription = "MFCC特征 (time_steps x features)"
    spec.description.output[0].shortDescription = "咀嚼概率 (0-1)"
    
    # 重新创建模型
    coreml_model = ct.models.MLModel(spec)
    
    # 保存模型
    print(f"\n保存CoreML模型: {output_path}")
    coreml_model.save(output_path)
    
    print("\n✅ 转换完成！")
    print(f"   CoreML模型: {output_path}")
    print(f"   输入形状: (1, {input_shape[0]}, {input_shape[1]})")
    print(f"   输出: chewing_probability (0-1)")
    
    return coreml_model

def test_coreml_model(model_path: str = "ChewingDetector.mlmodel"):
    """测试CoreML模型"""
    print(f"\n测试CoreML模型: {model_path}")
    
    # 加载
    model = ct.models.MLModel(model_path)
    
    # 创建测试输入
    test_input = {
        "input_features": np.random.randn(1, 1, 117).astype(np.float32)
    }
    
    # 预测
    prediction = model.predict(test_input)
    
    print(f"✅ 测试成功！")
    print(f"   输入形状: {test_input['input_features'].shape}")
    print(f"   输出: {prediction}")

if __name__ == "__main__":
    # 转换模型
    convert_to_coreml()
    
    # 测试
    test_coreml_model()
    
    print("\n📱 下一步:")
    print("1. 将生成的 ChewingDetector.mlmodel 拖入Xcode项目")
    print("2. 更新 RealtimeChewingDetector.swift 参数:")
    print("   - sampleRate: 25Hz")
    print("   - windowSize: 64")
    print("   - hopSize: 32")
    print("3. 实现MFCC特征提取")
